package algorithm.graph;

import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

/**
 * p 算法
 *
 * @author csp 2021-12-31
 */
public class COde05_Prime {

    static Comparator<Edge> edgeComparator = (o1, o2)->o2.weight - o1.weight;

    public static Set<Edge> primMST(Graph graph) {

        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(edgeComparator);

        HashSet<Node> set = new HashSet<>();
        HashSet<Edge> result = new HashSet<>();

        for (Node node : graph.nodes.values()) { // 随便挑了一个点

            // node 是开始点
            if (set.add(node)) {
                for (Edge edge : node.edges) {
                    priorityQueue.add(edge);
                }
                while (!priorityQueue.isEmpty()) {
                    Edge edge = priorityQueue.poll(); // 弹出解锁的边中最小的
                    Node toNode = edge.to; //可能的一个新的点
                    if (set.add(toNode)) { // 不含有的时候就是新的点
                        set.add(toNode);
                        result.add(edge);
                        for (Edge nextEdge : toNode.edges) {
                            priorityQueue.add(nextEdge);
                        }
                    }
                }
            }
        }
        return result;
    }

}
